//
// Created by bobin on 18-3-21.
//
#include "caffe/caffe.hpp"
#include "caffe/net.hpp"
#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "string"
using namespace std;
using namespace caffe;

const string data_root = "/home/bobin/data/kitti/odom/dataset/sequences/00/";
const string model_file = "/home/bobin/Documents/dnn/seg/PSPNet/evaluation/model/pspnet101_VOC2012.caffemodel";
const string model_deploy = "/home/bobin/Documents/dnn/seg/PSPNet/evaluation/prototxt/pspnet101_VOC2012_473.prototxt";
class PSP_Net{
public:
    PSP_Net(const std::string &deploy_file, const std::string &model_file);
    void forward(const cv::Mat& img);
    void WrapInputLayer(std::vector<cv::Mat>* input_channels);
    cv::Mat Preprocess(const cv::Mat im, std::vector<cv::Mat>* input_channels);
    enum DATA_TYPE {
        ADE20K = 0,
        VOC2012 = 1,
        cityscapes = 2
    };
    void config_data(DATA_TYPE data_type);
private:
    shared_ptr<caffe::Net<float> > net;
    bool isVal; //evaluation on valset
    int step; //equals to number of images divide num of GPUs in testing e.g. 500=2000/4
    int fea_cha; //number of classes
    int base_size; //based size for scaling
    int crop_size; //crop size fed into network
    string data_class; //class name
    string data_colormap; //color map
};

PSP_Net::PSP_Net(const std::string &deploy_file, const std::string &model_file) {
    Caffe::set_mode(Caffe::CPU);
    net.reset(new caffe::Net<float>(model_deploy, TEST));
    net->CopyTrainedLayersFrom(model_file);
}

void PSP_Net::config_data(DATA_TYPE data_type) {

    switch (data_type){
        case ADE20K:
            isVal = true; //evaluation on valset
            step = 500; //equals to number of images divide num of GPUs in testing e.g. 500=2000/4
            fea_cha = 150; //number of classes
            base_size = 512; //based size for scaling
            crop_size = 473; //crop size fed into network
            data_class = 'objectName150.mat'; //class name
            data_colormap = 'color150.mat'; //color map
            break;
        case VOC2012:
            isVal = false; //evaluation on testset
            step = 364; //364=1456/4
            fea_cha = 21;
            base_size = 512;
            crop_size = 473;
            data_class = 'objectName21.mat';
            data_colormap = 'colormapvoc.mat';
            break;
        case cityscapes:
            isVal = true;
            step = 125; //125=500/4
            fea_cha = 19;
            base_size = 2048;
            crop_size = 713;
            data_class = 'objectName19.mat';
            data_colormap = 'colormapcs.mat';
    }

}

void PSP_Net::WrapInputLayer(std::vector<cv::Mat>* input_channels) {
    Blob<float>* input_layer = net->input_blobs()[0];
    int width = input_layer->width();
    int height = input_layer->height();
    float* input_data = input_layer->mutable_cpu_data();
    for (int i = 0; i < input_layer->channels(); ++i) {
        cv::Mat channel(height, width, CV_32FC1, input_data);
        input_channels->push_back(channel);
        input_data += width * height;
    }
}


cv::Mat PSP_Net::Preprocess(const cv::Mat im, std::vector<cv::Mat>* input_channels) {
    int row = im.rows;
    int col = im.cols;
    float mean_r = 123.68; //means to be subtracted and the given values are used in our training stage
    float mean_g = 116.779;
    float mean_b = 103.939;

    cv::Mat im_pad, img_norm, im_mean;
    cv::Mat img_r_mean(row, col, CV_32F, mean_r);
    cv::Mat img_g_mean(row, col, CV_32F, mean_g);
    cv::Mat img_b_mean(row, col, CV_32F, mean_b);

    std::vector<cv::Mat> mean_channels;
    mean_channels.push_back(img_b_mean);
    mean_channels.push_back(img_g_mean);
    mean_channels.push_back(img_r_mean);
    cv::merge(mean_channels, im_mean);
    im.convertTo(im_pad, CV_32FC3);
    cv::subtract(im_pad, im_mean, img_norm);
    cv::split(img_norm, *input_channels);

}

void PSP_Net::forward(const cv::Mat &img) {
    std::vector<cv::Mat> input_channels;
    WrapInputLayer(&input_channels);

    Preprocess(img, &input_channels);
    net->ForwardPrefilled();
}

void LoadImages(const string &strPathToSequence, vector<string> &vstrImageLeft,
                vector<string> &vstrImageRight, vector<double> &vTimestamps)
{
    ifstream fTimes;
    string strPathTimeFile = strPathToSequence + "/times.txt";
    fTimes.open(strPathTimeFile.c_str());
    while(!fTimes.eof())
    {
        string s;
        getline(fTimes,s);
        if(!s.empty())
        {
            stringstream ss;
            ss << s;
            double t;
            ss >> t;
            vTimestamps.push_back(t);
        }
    }

    string strPrefixLeft = strPathToSequence + "/image_2/";
    string strPrefixRight = strPathToSequence + "/image_3/";

    const int nTimes = vTimestamps.size();
    vstrImageLeft.resize(nTimes);
    vstrImageRight.resize(nTimes);

    for(int i=0; i<nTimes; i++)
    {
        stringstream ss;
        ss << setfill('0') << setw(6) << i;
        vstrImageLeft[i] = strPrefixLeft + ss.str() + ".png";
        vstrImageRight[i] = strPrefixRight + ss.str() + ".png";
    }
}


int main(int argc, char**argv){
    vector<string> vstrImageLeft;
    vector<string> vstrImageRight;
    vector<double> vTimestamps;
    LoadImages(data_root, vstrImageLeft, vstrImageRight, vTimestamps);
    const int nImages = vstrImageLeft.size();
    PSP_Net psp_net(model_deploy, model_file);
    vector<float> vTimesTrack;
    vTimesTrack.resize(nImages);

    cout << endl << "-------" << endl;
    cout << "Start processing sequence ..." << endl;
    cout << "Images in the sequence: " << nImages << endl << endl;

    // Main loop
    cv::Mat imLeft, imRight;
    for(int ni=0; ni<nImages; ni++) {
        // Read left and right images from file
        imLeft = cv::imread(vstrImageLeft[ni], CV_LOAD_IMAGE_UNCHANGED);
        imRight = cv::imread(vstrImageRight[ni], CV_LOAD_IMAGE_UNCHANGED);
        double tframe = vTimestamps[ni];

        if (imLeft.empty()) {
            cerr << endl << "Failed to load image at: "
                 << string(vstrImageLeft[ni]) << endl;
            return 1;
        }
        cv::imshow("test", imLeft);
        cv::waitKey(10);
    }
}
